package storage

import (
	"database/sql"
	"fmt"
	_ "github.com/lib/pq"
)

type PostgresStorage struct {
	db *sql.DB
}

func NewPostgresStorage(connStr string) (*PostgresStorage, error) {
	db, err := sql.Open("postgres", connStr)
	if err != nil {
		return nil, fmt.Errorf("failed to connect to PostgreSQL: %w", err)
	}

	// Test connection
	if err := db.Ping(); err != nil {
		return nil, fmt.Errorf("failed to ping PostgreSQL: %w", err)
	}

	// Create table if not exists
	createTableSQL := `
	CREATE TABLE IF NOT EXISTS url_mappings (
		id SERIAL PRIMARY KEY,
		short_code VARCHAR(20) NOT NULL UNIQUE,
		long_url TEXT NOT NULL,
		created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
	);
	CREATE UNIQUE INDEX IF NOT EXISTS idx_short_code ON url_mappings(short_code);
	CREATE INDEX IF NOT EXISTS idx_long_url ON url_mappings(long_url);
	`

	_, err = db.Exec(createTableSQL)
	if err != nil {
		return nil, fmt.Errorf("failed to create table: %w", err)
	}

	return &PostgresStorage{db: db}, nil
}

func (s *PostgresStorage) Save(longURL, shortCode string) error {
	// 使用事务确保数据一致性
	tx, err := s.db.Begin()
	if err != nil {
		return err
	}
	defer tx.Rollback()

	// 检查长URL是否已存在
	var existingCode string
	err = tx.QueryRow("SELECT short_code FROM url_mappings WHERE long_url = $1", longURL).Scan(&existingCode)

	if err == nil {
		return ErrURLExists
	} else if err != sql.ErrNoRows {
		return err
	}

	// 插入新映射
	_, err = tx.Exec("INSERT INTO url_mappings (short_code, long_url) VALUES ($1, $2)",
		shortCode, longURL)
	if err != nil {
		return err
	}

	return tx.Commit()
}

func (s *PostgresStorage) Get(shortCode string) (string, error) {
	var longURL string
	err := s.db.QueryRow("SELECT long_url FROM url_mappings WHERE short_code = $1", shortCode).Scan(&longURL)

	if err == sql.ErrNoRows {
		return "", ErrURLNotFound
	}
	return longURL, err
}

func (s *PostgresStorage) Exists(longURL string) (string, bool) {
	var shortCode string
	err := s.db.QueryRow("SELECT short_code FROM url_mappings WHERE long_url = $1", longURL).Scan(&shortCode)

	if err != nil {
		return "", false
	}
	return shortCode, true
}

func (s *PostgresStorage) Close() error {
	return s.db.Close()
}
